{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# nvImageCodec with cv-cuda (Linux only)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from matplotlib import pyplot as plt\n",
    "import cupy as cp\n",
    "\n",
    "import cvcuda"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Setting resource folder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "resources_dir = os.getenv(\"PYNVIMGCODEC_EXAMPLES_RESOURCES_DIR\", \"../assets/images/\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Import nvImageCodec module and create Decoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nvidia import nvimgcodec\n",
    "decoder = nvimgcodec.Decoder()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Read image with nvImageCodec"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputImage = decoder.read(resources_dir + \"tabby_tiger_cat.jpg\")\n",
    "print(\"size:{}x{}\".format(inputImage.width, inputImage.height))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Pass it to cvcuda using as_tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nvcvInputTensor = cvcuda.as_tensor(inputImage, \"HWC\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Resize with cvcuda"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cvcuda_stream = cvcuda.Stream()\n",
    "with cvcuda_stream:\n",
    "    nvcvResizeTensor = cvcuda.resize(nvcvInputTensor, (320, 320, 3), cvcuda.Interp.CUBIC)\n",
    "    nvcvResizeTensor.cuda().__cuda_array_interface__"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Write with nvImageCodec"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "encoder = nvimgcodec.Encoder()\n",
    "encoder.write(\"tabby_tiger_cat_320x320.jpg\", nvimgcodec.as_image(nvcvResizeTensor.cuda(), cuda_stream = cvcuda_stream.handle))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Verify with OpenCV"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "image = cv2.imread(\"tabby_tiger_cat_320x320.jpg\")\n",
    "image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
    "plt.imshow(image)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Resnet50 classification example from cv-cuda"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torchvision.io.image import read_file, decode_jpeg\n",
    "from torchvision import models\n",
    "import numpy as np\n",
    "file_name = resources_dir + \"tabby_tiger_cat.jpg\"\n",
    "labelsfile = resources_dir + \"../imagenet-classes.txt\"\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Orignal example code uses torchvision to load image (under the hood it uses nvJpeg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = read_file(file_name)\n",
    "inputImageTmp = decode_jpeg(data, device=\"cuda\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can change this code to use nvImageCodec and use all formats available with plugins. Please uncomment lines with other images to test it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#file_name = resources_dir + \"cat-1046544_640.jp2\"\n",
    "#file_name = resources_dir + \"Weimaraner.jpg\"\n",
    "\n",
    "# nvImageCodec has fallback for cpu decoder (only when necessary plugins are installed)\n",
    "# for codec do not supported yet on GPU so we can read e.g. png \n",
    "#file_name = resources_dir + \"cat-1245673_640.png\"\n",
    "\n",
    "inputImage = decoder.read(file_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# A torch tensor/ or nvImageCodec Image can be wrapped into a CVCUDA Object using the \"as_tensor\"\n",
    "# function in the specified layout. The datatype and dimensions are derived\n",
    "# directly from the torch tensor.\n",
    "nvcvInputTensor = cvcuda.as_tensor(inputImage, \"HWC\")\n",
    "image = cp.asnumpy(nvcvInputTensor.cuda())\n",
    "plt.imshow(image)\n",
    "\n",
    "#Need 4 dimensions when first is batch size\n",
    "image_tensors = torch.stack((torch.as_tensor(nvcvInputTensor.cuda()),))\n",
    "nvcvInputTensor = cvcuda.as_tensor(image_tensors.cuda(), \"NHWC\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Preprocessing includes the following sequence of operations.\n",
    "Resize -> DataType Convert(U8->F32) -> Normalize(Apply mean and std deviation)\n",
    "-> Interleaved to Planar\n",
    "\"\"\"\n",
    "\n",
    "# Model settings\n",
    "layerHeight = 224\n",
    "layerWidth = 224\n",
    "batchSize = 1\n",
    "\n",
    "# Resize\n",
    "# Resize to the input network dimensions\n",
    "nvcvResizeTensor = cvcuda.resize(nvcvInputTensor, (1, layerHeight, layerWidth, 3), cvcuda.Interp.CUBIC)\n",
    "    \n",
    "# Convert to the data type and range of values needed by the input layer\n",
    "# i.e uint8->float. A Scale is applied to normalize the values in the range 0-1\n",
    "nvcvConvertTensor = cvcuda.convertto(nvcvResizeTensor, np.float32, scale=1 / 255)\n",
    "\n",
    "\"\"\"\n",
    "The input to the network needs to be normalized based on the mean and\n",
    "std deviation value to standardize the input data.\n",
    "\"\"\"\n",
    "\n",
    "# Create a torch tensor to store the mean and standard deviation values for R,G,B\n",
    "scale = [0.485, 0.456, 0.406]\n",
    "std = [0.229, 0.224, 0.225]\n",
    "scaleTensor = torch.Tensor(scale)\n",
    "stdTensor = torch.Tensor(std)\n",
    "\n",
    "# Reshape the the number of channels. The R,G,B values scale and offset will be\n",
    "# applied to every color plane respectively across the batch\n",
    "scaleTensor = torch.reshape(scaleTensor, (1, 1, 1, 3)).cuda()\n",
    "stdTensor = torch.reshape(stdTensor, (1, 1, 1, 3)).cuda()\n",
    "\n",
    "# Wrap the torch tensor in a CVCUDA Tensor\n",
    "nvcvScaleTensor = cvcuda.as_tensor(scaleTensor, \"NHWC\")\n",
    "nvcvBaseTensor = cvcuda.as_tensor(stdTensor, \"NHWC\")\n",
    "\n",
    "# Apply the normalize operator and indicate the scale values are std deviation\n",
    "# i.e scale = 1/stddev\n",
    "nvcvNormTensor = cvcuda.normalize(nvcvConvertTensor,\n",
    "    nvcvBaseTensor, nvcvScaleTensor, cvcuda.NormalizeFlags.SCALE_IS_STDDEV\n",
    ")\n",
    "\n",
    "# The final stage in the preprocess pipeline includes converting the RGB buffer\n",
    "# into a planar buffer\n",
    "nvcvPreprocessedTensor = cvcuda.reformat(nvcvNormTensor, \"NCHW\")\n",
    "\n",
    "# Inference uses pytorch to run a resnet50 model on the preprocessed input and outputs\n",
    "# the classification scores for 1000 classes\n",
    "# Load Resnet model pretrained on Imagenet\n",
    "resnet50 = models.resnet50(pretrained=True)\n",
    "resnet50.to(\"cuda\")\n",
    "resnet50.eval()\n",
    "\n",
    "# Run inference on the preprocessed input\n",
    "torchPreprocessedTensor = torch.as_tensor(nvcvPreprocessedTensor.cuda(), device=\"cuda\")\n",
    "inferOutput = resnet50(torchPreprocessedTensor)\n",
    "\n",
    "\"\"\"\n",
    "Postprocessing function normalizes the classification score from the network and sorts\n",
    "the scores to get the TopN classification scores.\n",
    "\"\"\"\n",
    "# top results to print out\n",
    "topN = 5\n",
    "\n",
    "# Read and parse the classes\n",
    "with open(labelsfile, \"r\") as f:\n",
    "    classes = [line.strip() for line in f.readlines()]\n",
    "\n",
    "# Apply softmax to Normalize scores between 0-1\n",
    "scores = torch.nn.functional.softmax(inferOutput, dim=1)[0]\n",
    "\n",
    "# Sort output scores in descending order\n",
    "_, indices = torch.sort(inferOutput, descending=True)\n",
    "\n",
    "# Display Top N Results\n",
    "for idx in indices[0][:topN]:\n",
    "    print(\"Class : \", classes[idx], \" Score : \", scores[idx].item())\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.10 64-bit",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "916dbcbb3f70747c44a77c7bcd40155683ae19c65e1c03b4aa3499c5328201f1"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
